import torch
import torch.nn as nn
import torch.nn.functional as nn_functional

from setting import batch_size
import torchinfo


class Model(nn.Module):

    def __init__(self):
        super(Model, self).__init__()

        # 完全由全连接成构成
        self.fc = nn.Sequential(
            nn.Flatten(),
            # 隐藏层
            nn.Linear(in_features=784, out_features=1024),
            nn.ReLU(),
            # 输出层
            nn.Linear(in_features=1024, out_features=10),
            # nn.Softmax(dim=1)
        )

    def forward(self, x):
        x = self.fc(x)
        # x = torch.argmax(x, dim=1)
        return x


def getModel():
    return Model()


if __name__ == '__main__':
    model = getModel()
    out = model(torch.randn(batch_size, 1, 28, 28))
    # print(out)
    # print(torchinfo.summary(model, (batch_size, 1, 28, 28)))
